# -*- coding: utf-8 -*-
"""
@Env 
@Time 2024/9/3 上午10:54
@Author yzpang
@Function: 模型管理接口
"""
import traceback

from ..entity.response import BaseResponse
from ..entity.request import BaseRequest, ModelMessage
from ..service import model_service
from modelserver.configs.base_config import get_logger

logger = get_logger()


def mount_model_routes(app):
    """
    挂载模型相关接口
    :param app:
    """

    app.post(path="/model/list",
             tags=['模型管理接口'],
             summary='列出当前已加载的模型',
             )(list_models)

    app.post(path="/model/load",
             tags=['模型管理接口'],
             summary="加载模型",
             )(load_model)

    app.post(path="/model/unload",
             tags=['模型管理接口'],
             summary='卸载模型',
             )(unload_model)


def list_models(request: BaseRequest) -> BaseResponse:
    """ 列出所有的模型列表"""
    logger.info(request)
    models = model_service.list_models()
    return BaseResponse(data=models)


def list_running_models(request: BaseRequest) -> BaseResponse:
    """ 列出运行的模型列表 """
    pass


def load_model(model: ModelMessage) -> BaseResponse:
    """ 载入模型 """
    logger.info(model)
    try:
        resp = model_service.load_model(model=model)
        return BaseResponse(data=resp)
    except Exception as e:
        list_models(BaseRequest())
        traceback.print_exc()
        return BaseResponse(code=500, msg="failed", data=str(e))


def unload_model(model: ModelMessage) -> BaseResponse:
    """ 卸载模型 """
    logger.info(model)
    try:
        resp = model_service.unload_model(model=model)
        return BaseResponse(data=resp)
    except Exception as e:
        traceback.print_exc()
        return BaseResponse(code=500, msg="failed", data=str(e))



